#include "Error.h"
#include "ErrorUtils.h"
#include "RegistryManager.h"
#include "StringUtils.h"

namespace
{
    const std::wstring HKCR(L"HKEY_CLASSES_ROOT");
    const std::wstring HKCU(L"HKEY_CURRENT_USER");
    const std::wstring HKLM(L"HKEY_LOCAL_MACHINE");
    const std::wstring HKU(L"HKEY_USERS");
    const std::wstring HKCC(L"HKEY_CURRENT_CONFIG");

    TestFramework::TestWString ExtractKeyName(const TestFramework::TestWString& path)
    {
        TestFramework::TestWString test(L"\\");
        std::size_t index = path.RFind(test);
        if (index == path.NPos)
        {
            return path;
        }

        return TestFramework::TestWString(path.Begin() + index + 1, path.End());
    }

    template <class T>
    std::size_t GetCharsCount(const TestFramework::TestVector<T>& data)
    {
        std::size_t size = 0;
        TestFramework::TestVector<T>::ConstIterator iter = data.Begin();
        for (; iter != data.End(); ++iter)
        {
            size += iter->GetLength();
        }

        return size;
    }

    void CheckResultAndThrow(LSTATUS result, const char* errorDescription, int errorCode)
    {
        if (result != ERROR_SUCCESS)
        {
            TestFramework::TestString errorText(errorDescription);
            errorText += " Error code is \'";
            errorText += TestFramework::GetErrorCodePattern();
            errorText += "'.";
            throw std::exception(TestFramework::FormatErrorText(errorText.Data(), result), errorCode);
        }
    }

    void ApplyData(const HKEY& key, const TestFramework::TestWString& name, const void* data,
                   const int size, const DWORD dataType)
    {
        LSTATUS result = ::RegSetValueEx(key, name.Data(), 0, dataType, (const byte*)data, size);
        CheckResultAndThrow(result, "Failed to write registry value.",
                            TestFramework::Errors::FAILED_TO_WRITE_REGISTRY_VALUE);
    }

    template <typename T>
    void GetNumericValue(const HKEY& key, const TestFramework::TestWString& name, T& data, int size)
    {
        LSTATUS result = ::RegQueryValueEx(key, name.Data(), 0, 0, (LPBYTE)&data, (LPDWORD)&size);
        CheckResultAndThrow(result, "Failed to get registry value.",
                            TestFramework::Errors::FAILED_TO_GET_REGISTRY_VALUE);
    }
}

namespace TestFramework
{
    RegistryValue::RegistryValue(const HKEY& root, const TestWString& name)
        : Key(root)
        , Name(name)
    {
    }

    int RegistryValue::GetValueSize() const
    {
        std::size_t size = 0;
        LSTATUS result = ::RegQueryValueEx(Key, Name.Data(), 0, 0, 0, (LPDWORD)&size);
        CheckResultAndThrow(result, "Failed to get registry value size.", Errors::FAILED_TO_GET_REGISTRY_VALUE_INFO);
        return size;
    }

    void RegistryValue::SetValue(const TestVector<byte>& data)
    {
        ApplyData(Key, Name, &data[0], static_cast<int>(data.GetSize()), REG_BINARY);
    }

    void RegistryValue::SetValue(const DWORD data)
    {
        ApplyData(Key, Name, &data, sizeof(DWORD), REG_DWORD);
    }

    void RegistryValue::SetValue(const TestWString& data)
    {
        ApplyData(Key, Name, data.Data(), static_cast<int>(sizeof(TestWString::CharType) * data.GetLength()), REG_SZ);
    }

    void RegistryValue::SetValue(const long long data)
    {
        ApplyData(Key, Name, &data, sizeof(long long), REG_QWORD);
    }

    void RegistryValue::GetValue(TestVector<byte>& data, int size) const
    {
        byte* buff = new byte[size]();
        LSTATUS result = ::RegQueryValueEx(Key, Name.Data(), 0, 0, (LPBYTE)buff, (LPDWORD)&size);
        CheckResultAndThrow(result, "Failed to get registry value.", Errors::FAILED_TO_GET_REGISTRY_VALUE);

        data.Clear();
        data.Reserve(size);
        for (int i = 0; i < size; ++i)
        {
            data.PushBack(*((byte*)buff + i));
        }
    }

    void RegistryValue::GetValue(DWORD& data, int size) const
    {
        GetNumericValue(Key, Name.Data(), data, size);
    }

    void RegistryValue::GetValue(TestWString& data, int size) const
    {
        TestWString::CharType* buff = new TestWString::CharType[size / sizeof(TestWString::CharType)];
        LSTATUS result = ::RegQueryValueEx(Key, Name.Data(), 0, 0, (LPBYTE)buff, (LPDWORD)&size);
        CheckResultAndThrow(result, "Failed to get registry value.", Errors::FAILED_TO_GET_REGISTRY_VALUE);
        data.Clear();
        data += buff;
    }

    void RegistryValue::GetValue(long long& data, int size) const
    {
        GetNumericValue(Key, Name.Data(), data, size);
    }

    ///////////////////////////////////////////////////////////////////////////////////////////////

    RegistryKey::RegistryKey(HKEY root, int accessType)
        : Key(root)
        , Access(accessType)
    {
    }

    RegistryKey::RegistryKey(HKEY root, const TestWString& path, int accessType)
        : Key(root)
        , Access(accessType)
    {
        LSTATUS result = ::RegOpenKeyExW(Key, path.Data(), 0, Access, &Key);
        CheckResultAndThrow(result, "Failed to open registry key.", Errors::FAILED_TO_OPEN_REGISTRY_KEY);
    }

    RegistryKey::~RegistryKey()
    {
        if (Key && !Name.IsEmpty())
        {
            ::RegCloseKey(Key);
        }
    }

    void RegistryKey::Close()
    {
        if (Key && !Name.IsEmpty())
        {
            ::RegCloseKey(Key);
        }
    }

    RegistryKey RegistryKey::operator[](const TestWString& path)
    {
        Name = ExtractKeyName(path);
        HKEY tmp;
        LSTATUS result = ::RegOpenKeyExW(Key, path.Data(), 0, Access, &tmp);
        CheckResultAndThrow(result, "Failed to open registry key.", Errors::FAILED_TO_OPEN_REGISTRY_KEY);
        return RegistryKey(tmp, Access);
    }

    RegistryKey RegistryKey::OpenOrCreateKey(const TestWString& path)
    {
        Name = ExtractKeyName(path);
        HKEY tmp;
        LSTATUS result = ::RegOpenKeyExW(Key, path.Data(), 0, Access, &tmp);
        if (result != ERROR_SUCCESS)
        {
            return CreateKey(path);
        }

        return RegistryKey(tmp, Access);
    }

    RegistryKey RegistryKey::CreateKey(const TestWString& path)
    {
        HKEY tmp;
        unsigned regRes;
        LSTATUS result = ::RegCreateKeyEx(Key, path.Data(), 0, 0, REG_OPTION_VOLATILE,
                                          Access, 0, &tmp, (LPDWORD)&regRes);

        CheckResultAndThrow(result, "Failed to create registry key.", Errors::FAILED_TO_CREATE_REGISTRY_KEY);
        return RegistryKey(tmp, Access);
    }

    HKEY RegistryKey::operator()() const
    {
        return Key;
    }

    KeyInfo RegistryKey::GetInfo() const
    {
        KeyInfo info;
        LSTATUS result = ::RegQueryInfoKey(Key, 0, 0, 0, &info.SubKeysCount, 0, 0, &info.ValuesCount, 0, 0,
            &info.SecurityDescrptor, &info.LastWrite);

        CheckResultAndThrow(result, "Failed to get registry key info.", Errors::FAILED_TO_GET_REGISTRY_INFO);
        return info;
    }

    TestAutoPtr<RegistryValue> RegistryKey::CreateValue(const TestWString& name) const
    {
        return TestAutoPtr<RegistryValue>(new RegistryValue(Key, name));
    }

    TestAutoPtr<RegistryValue> RegistryKey::OpenValue(const TestWString& name) const
    {
        return TestAutoPtr<RegistryValue>(new RegistryValue(Key, name));
    }

    void RegistryKey::DeleteValue(const TestWString& name) const
    {
        LSTATUS result = ::RegDeleteValue(Key, name.Data());
        CheckResultAndThrow(result, "Failed to delete registry value.", Errors::FAILED_TO_DELETE_REGISTRY_VALUE);
    }

    void RegistryKey::DeleteKey(const TestWString& path) const
    {
        LSTATUS result = ::RegDeleteKeyEx(Key, path.Data(), Access, 0);
        CheckResultAndThrow(result, "Failed to delete registry key.", Errors::FAILED_TO_DELETE_REGISTRY_KEY);
    }

    void RegistryKey::EnumerateKeys(KeysEnumerator& processor) const
    {
        KeyInfo info = GetInfo();
        wchar_t buf[1024] = {0};
        
        for (DWORD i = 0; i < info.SubKeysCount; ++i)
        {
            DWORD bufLength = sizeof(buf) / sizeof(buf[0]);
            LSTATUS result = ::RegEnumKeyEx(Key, i, buf, &bufLength, 0, 0, 0, 0);
            CheckResultAndThrow(result, "Failed to enumerate registry key.", Errors::FAILED_TO_ENUMERATE_REGISTRY_KEY);
            TestWString name(buf);
            RegistryKey item(Key, name, Access);
            processor.ProcessKey(item);
        }
    }

    void RegistryKey::EnumerateValues(ValueEnumerator& processor) const
    {
        KeyInfo info = GetInfo();
        wchar_t buf[1024];
        
        for (DWORD i = 0; i < info.ValuesCount; ++i)
        {
            DWORD bufLength = sizeof(buf) / sizeof(buf[0]);
            RegEnumValue(Key, i, buf, &bufLength, 0, 0, 0, 0);
            TestWString name(buf);
            TestAutoPtr<RegistryValue> value(OpenValue(name));
            processor.ProcessValue(*value);
        }
    }

    ///////////////////////////////////////////////////////////////////////////////////////////////

    RegistryKey GetClassesRoot(int accessType)
    {
        return RegistryKey(HKEY_CLASSES_ROOT, accessType);
    }

    RegistryKey GetCurrentUser(int accessType)
    {
        return RegistryKey(HKEY_CURRENT_USER, accessType);
    }

    RegistryKey GetLocalMachine(int accessType)
    {
        return RegistryKey(HKEY_LOCAL_MACHINE, accessType);
    }

    RegistryKey GetUsers(int accessType)
    {
        return RegistryKey(HKEY_USERS, accessType);
    }

    RegistryKey GetCurrentConfig(int accessType)
    {
        return RegistryKey(HKEY_CURRENT_CONFIG, accessType);
    }
}
